import numpy as np

from algorithms.pattern_recognition import BandStrategy, ShortTermStrategy
from data.data_fetcher import DataFetcher
from data.data_processor import DataProcessor
from data.data_visualizer import DataVisualizer
from database.database_manager import DatabaseManager
from execution.trade_executor import TradeDetector
from config.load_config import load_strategy_config
import tushare as ts
import pymysql
from datetime import datetime, timedelta
import openpyxl
def main():
    # 初始化相关组件和对象
    # global detect_date
    global detect_date
    token = 'b6d9a5208e4d4b30715424b58094e8cd89a80afae4a76e6eb36727f0'
    database = 'manhattan_db'
    pro = ts.pro_api()
    all_stock_name = pro.stock_basic(exchange='', list_status='L', fields='ts_code, name')
    stock_code_name = {}  # 代码与名称的映射
    stock_pool = []
    for i in range(all_stock_name['ts_code'].shape[0]):
        code = all_stock_name['ts_code'][i]
        if '*ST' not in code:
            stock_pool.append(code)
            stock_code_name[code] = all_stock_name['name'][i]
    print("select ", len(stock_pool), "stocks!")
    processed_data = np.array([])

    # 加载策略配置
    strategy_config = load_strategy_config('config/strategy_config.yaml')
    strategy_type = strategy_config['strategy_type']
    save_path = strategy_config['vis_save_path']
    download_from_tushare = strategy_config['download_from_tushare']
    start_date = strategy_config['start_date']
    end_date = strategy_config['end_date']
    # detect_date = strategy_config['detect_date']
    detect_start_date = strategy_config['detect_start_date']
    detect_end_date = strategy_config['detect_end_date']
    resort_excel = True
    # 初始化模块
    data_fetcher = DataFetcher(token)
    data_processor = DataProcessor()
    database_manager = DatabaseManager(database)
    data_visualizer = DataVisualizer(save_path=save_path, visual_days=60)
    trade_executor = TradeDetector()

    if download_from_tushare == 'daily_stock':
        print("Download From Tushare")
        import time
        start_time = time.time()
        # 获取市场数据
        market_data = data_fetcher.fetch_data(symbol=stock_pool, start_date=start_date,
                                              end_date=end_date, data_type='daily')
        block1_time = time.time() - start_time

        # 处理市场数据
        start_time = time.time()
        # processed_data = data_processor.process_data(market_data)
        processed_data = market_data
        block2_time = time.time() - start_time

        # 存储数据到数据库
        start_time = time.time()
        database_manager.store_data(processed_data)
        block3_time = time.time() - start_time
        print(block1_time, block2_time, block3_time)
    elif download_from_tushare == 'daily_turnover':
        print("Download turnover From Tushare")
        turnover_data = data_fetcher.fetch_turnover(symbol=stock_pool[:10],
                                                    start_date_string='20230601', end_date_string='20230728')
        database_manager.store_turnover(turnover_data)

    else:  # download from database
        start_date_obj = datetime.strptime(detect_start_date, "%Y%m%d")
        end_date_obj = datetime.strptime(detect_end_date, "%Y%m%d")
        cur_detect_date = start_date_obj

        while cur_detect_date <= end_date_obj:
            database_manager = DatabaseManager(database)
            weekday_num = cur_detect_date.weekday()
            if database == "mydatabase":
                processed_data = database_manager.fetch_data(table='stock_daily', turnover_table='turnover_daily',
                                                             start_date=start_date, end_date=end_date)
            if database == "manhattan_db":
                detect_date = cur_detect_date.strftime("%Y%m%d")
                processed_data = database_manager.fetch_manhattan(start_date=start_date, end_date=end_date,
                                                                  detect_date=detect_date)
            if 0 <= weekday_num <= 4:
                # 选择并初始化策略
                strategy = ''
                if strategy_type == 'short_term':
                    strategy = ShortTermStrategy(strategy_config, detect_date)
                if strategy_type == 'band':
                    strategy = BandStrategy(strategy_config)

                # 执行策略，筛选出符合特征的name，以及区间
                selected_stocks = strategy.execute_strategy(processed_data)

                # 可视化区间后某段时间并存储在excel
                data_visualizer.visualize_data(selected_stocks, processed_data, stock_code_name, detect_date)

                # 在后期每天都要检测出需要提醒的股票代码
                current_date = '20230718'
                trade_executor.detection(selected_stocks, current_date)
                # trade_executor.execute_trades(trade_signals)
            cur_detect_date += timedelta(days=1)  # 增加一天

if __name__ == "__main__":
    main()